Codes

# Function to generate the record count matrix for a single hospital
generate_record_count <- function(data) {
  counts <- table(data[, "n_hi"])
  result_matrix <- cbind(as.numeric(names(counts)), as.numeric(counts))
  colnames(result_matrix) <- c("n_hi", "frequency")
  return(result_matrix)
}

# Function to generate Z_hv matrix for a single hospital
generate_Zhv_matrix <- function(data) {
  record_count_matrix <- generate_record_count(data)
  diagonal_blocks <- list()

  for (i in 1:nrow(record_count_matrix)) {
    n_hi <- record_count_matrix[i, "n_hi"]
    frequency <- record_count_matrix[i, "frequency"]

    identity_block <- diag(frequency)
    ones_vector <- matrix(1, nrow = n_hi, ncol = 1)

    kronecker_product <- kronecker(identity_block, ones_vector)
    diagonal_blocks[[i]] <- kronecker_product
  }

  big_matrix <- do.call(Matrix::bdiag, diagonal_blocks)
  big_matrix <- cbind(1, big_matrix)
  return(as.matrix(big_matrix))
}

# Function to get summary stats from each site for distributed LMM
lmm.get.summary3 <- function(Y = NULL, X = NULL, Z = NULL, id.site = NULL, weights = NULL) {
  if (is.null(weights)) weights <- rep(1, length(Y))
  X <- as.matrix(X)
  id.site <- as.character(id.site)
  id.site.uniq <- unique(id.site)
  px <- ncol(X)

  ShXYZ <- list()
  for (h in seq_along(id.site.uniq)) {
    sh <- id.site.uniq[h]
    wth <- weights[id.site == sh]
    Xh <- X[id.site == sh, ]
    Yh <- Y[id.site == sh]
    Zh <- Z[h][[1]]

    ShX  <- t(Xh * wth) %*% Xh
    ShXZ <- t(Xh * wth) %*% Zh
    ShXY <- t(Xh * wth) %*% Yh
    ShZ  <- t(Zh * wth) %*% Zh
    ShZY <- t(Zh * wth) %*% Yh
    ShY  <- sum(Yh^2 * wth)
    Nh <- sum(id.site == sh)

    ShXYZ[[sh]] <- list(ShX = ShX, ShXZ = ShXZ, ShXY = ShXY,
                        ShZ = ShZ, ShZY = ShZY, ShY = ShY, Nh = Nh)
  }

  return(ShXYZ)
}

##############################
# different site different patients

# Parameters
H <- 5  # of sites
m_hosp <- sample(100:150, H, replace = T) # of patients (1k to 3k later)


px <- 9  # of covariates
p_bin <- 5  # of binary X
p_cont <- px - p_bin  # of continuous X


beta0 = 5 # intercept
beta <- c(2, 4, 6, 8, 10, 3, 5, 7, 9)  # Fixed effects for covariates

sigma_e <- 1  # error variance
sigma_u <- 3 # site-level variance
sigma_v_hosp <- runif(H, min = 1, max = 5)  # Varying sigma_v by hospital


# Generate data
nn <- rep(m_hosp, times = 1)  # Number of patients per hospital
id.hosp <- rep(1:H, times = m_hosp)   # Hospital ID
id.pat <- sequence(nn)                # Patient ID
n_visits <- sample(1:30, sum(nn), replace = TRUE)  # number of visits of patients

# Expand hospital and patient IDs for visits
id.visit <- sequence(n_visits)
id.hosp.expanded <- rep(id.hosp, times = n_visits)
id.pat.expanded <- rep(id.pat, times = n_visits)

# Random effects
u_h <- rnorm(H, mean = 0, sd = sigma_u)  # Hospital effects
v_hi <- rnorm(sum(nn), mean = 0, sd = rep(sigma_v_hosp, times = m_hosp))  # Patient effects varying by hospital

# Expansion of hospital effects
u_h_patient <- rep(u_h, times = m_hosp)
u_h_expanded <- rep(u_h_patient, times = n_visits)

# Expansion of patient effects
v_hi_expanded <- rep(v_hi, times = n_visits)

# covariates
X_bin <- matrix(rbinom(sum(n_visits) * p_bin, size = 1, prob = 0.3), nrow = sum(n_visits), ncol = p_bin)
X_cont <- matrix(rnorm(sum(n_visits) * p_cont, mean = 0, sd = 1), nrow = sum(n_visits), ncol = p_cont)
X_hij <- cbind(X_bin, X_cont)  # Combine binary & continuous covariates

epsilon_hij <- rnorm(sum(n_visits), mean = 0, sd = sigma_e)

# Compute outcome
y_hij <- beta0 + X_hij %*% beta + u_h_expanded + v_hi_expanded + epsilon_hij

three_lvl_dat <- data.table(
  site = id.hosp.expanded,
  patient = id.pat.expanded,
  visit = id.visit,
  X_hij,
  Y = y_hij
) %>% data.frame()

setnames(three_lvl_dat, c("site", "patient", "visit", paste0("X", 1:px), "Y"))


# Preprocessing
# Calculate the number of visits per patient per site
visit_count <- three_lvl_dat %>%
  dplyr::group_by(site, patient) %>%
  dplyr::summarise(total_visits = n(), .groups = "drop")

# Reorder data (as a preprocess probably)
rearranged_data <- merge(three_lvl_dat, visit_count, by = c("site", "patient")) %>%
  arrange(site, total_visits, patient) %>%
  mutate(site = factor(site))


# XYZ
Y <- rearranged_data$Y
X <- as.matrix(rearranged_data[, paste0("X", 1:px)])
X <- cbind(1, X)
Z <- list()

for(i in 1:H){
  count_mat = rearranged_data %>%
                    filter(site == i) %>%
                    group_by(site, patient) %>%
                    dplyr::summarise(n_hi = n(), .groups = 'drop')

  Z[[i]] <- (generate_Zhv_matrix(count_mat))
}

id.site <- rearranged_data$site

ShXYZ <- lmm.get.summary3(Y, X, Z, id.site)

source("DLMM_engine3RI.R")
bench::mark(lmm.fit3(Y = NULL, X = NULL, Z = NULL,
                     id.site = NULL, weights = NULL,
                     pooled = F, reml = T,
                     common.s2 = T,
                     ShXYZ = ShXYZ,
                     corstr = 'independence',
                     mypar.init = NULL))

default mypar.init (var comp) =  1 1 1 1 1 1 
Normal exit from bobyqa  The number of function evaluations used is  1739default mypar.init (var comp) =  1 1 1 1 1 1 
Normal exit from bobyqa  The number of function evaluations used is  1739

Warning: Some expressions had a GC in every iteration; so filtering is disabled.

fit03.dlmm = lmm.fit3(Y = NULL, X = NULL, Z = NULL,
                      id.site = NULL, weights = NULL,
                      pooled = F, reml = T,
                      common.s2 = T,
                      ShXYZ = ShXYZ,  # only need summary stats
                      corstr = 'independence',
                      mypar.init = NULL)

default mypar.init (var comp) =  1 1 1 1 1 1 
Normal exit from bobyqa  The number of function evaluations used is  1739

fit03.dlmm$b

       [,1]
   2.427526
X1 2.030760
X2 4.025651
X3 6.009985
X4 8.027462
X5 9.964476
X6 3.002865
X7 5.003628
X8 7.003889
X9 8.990027

sqrt(fit03.dlmm$V)

[1] 2.006603 2.837770 3.293975 2.701197 2.025338 2.754136

sqrt(fit03.dlmm$s2)

[1] 0.9975977

true_beta = c(beta0, beta)
true_beta

 [1]  5  2  4  6  8 10  3  5  7  9

true_sigma = c(sigma_u, sigma_v_hosp)
true_sigma

[1] 3.000000 2.733776 3.263261 2.339320 2.124376 2.726359

c(sigma_e)

[1] 1

plot(fit03.dlmm$b, true_beta)
points(fit03.dlmm$b[1], true_beta[1], col = "red")

abline(a = 0, b = 1, col = "blue", lwd = 2)

plot(sqrt(fit03.dlmm$V), true_sigma)
points(sqrt(fit03.dlmm$V)[1], true_sigma[1], col = "red")

abline(a = 0, b = 1, col = "blue", lwd = 2)

source("DLMM_Engine.R")
bench::mark(lmm.fit3(Y = NULL, X = NULL, Z = NULL,
                     id.site = NULL, weights = NULL,
                     pooled = F, reml = T,
                     common.s2 = T,
                     ShXYZ = ShXYZ,
                     corstr = 'independence',
                     mypar.init = NULL))

Default mypar.init (var comp) =  1 1 1 1 1 1 
Normal exit from bobyqa  The number of function evaluations used is  1182 
Default mypar.init (var comp) =  1 1 1 1 1 1 
Normal exit from bobyqa  The number of function evaluations used is  1182 

Warning: Some expressions had a GC in every iteration; so filtering is disabled.

fit03.dlmm = lmm.fit3(Y = NULL, X = NULL, Z = NULL,
                      id.site = NULL, weights = NULL,
                      pooled = F, reml = T,
                      common.s2 = T,
                      ShXYZ = ShXYZ,  # only need summary stats
                      corstr = 'independence',
                      mypar.init = NULL)

Default mypar.init (var comp) =  1 1 1 1 1 1 
Normal exit from bobyqa  The number of function evaluations used is  1182 

fit03.dlmm$b

       [,1]
   2.427516
X1 2.030760
X2 4.025651
X3 6.009985
X4 8.027462
X5 9.964476
X6 3.002865
X7 5.003628
X8 7.003889
X9 8.990027

sqrt(fit03.dlmm$V)

[1] 2.003896 2.837504 3.294069 2.701259 2.025409 2.753977

sqrt(fit03.dlmm$s2)

[1] 0.9975991

true_beta = c(beta0, beta)
true_beta

 [1]  5  2  4  6  8 10  3  5  7  9

true_sigma = c(sigma_u, sigma_v_hosp)
true_sigma

[1] 3.000000 2.733776 3.263261 2.339320 2.124376 2.726359

c(sigma_e)

[1] 1

plot(fit03.dlmm$b, true_beta)
points(fit03.dlmm$b[1], true_beta[1], col = "red")

abline(a = 0, b = 1, col = "blue", lwd = 2)

plot(sqrt(fit03.dlmm$V), true_sigma)
points(sqrt(fit03.dlmm$V)[1], true_sigma[1], col = "red")

abline(a = 0, b = 1, col = "blue", lwd = 2)

source("DLMM_Engine_Efficient.R")
bench::mark(lmm.fit3(Y = NULL, X = NULL, Z = NULL,
                     id.site = NULL, weights = NULL,
                     pooled = F, reml = T,
                     common.s2 = T,
                     ShXYZ = ShXYZ,
                     corstr = 'independence',
                     mypar.init = NULL))

Default mypar.init (var comp) =  1 1 1 1 1 1 

A <- matrix(runif(1024^2), nrow = 1024)
B <- matrix(runif(1024^2), nrow = 1024)
size = 1024

bench::mark(logdet <- log(det(diag(1, size) + t(A) %*% A %*% B)))
bench::mark(logdet <- as.numeric(determinant(diag(1, size) + t(A) %*% A %*% B, logarithm = TRUE)$modulus))